#include <google/protobuf/message.h>
#include <zlib.h>

#include "zrpc/rpc/protobuf_codec_lite.h"
#include "zrpc/rpc/buffer_stream.h"
#include "zrpc/base/logger.h"
#include "zrpc/rpc/google.h"
#include "zrpc/net/tcp_connection.h"

using namespace zrpc;
namespace {
int ProtobufVersionCheck() {
  GOOGLE_PROTOBUF_VERIFY_VERSION;
  return 0;
}
int __attribute__((unused)) dummy = ProtobufVersionCheck();
}  // namespace

void ProtobufCodecLite::Send(const std::shared_ptr<TcpConnection>& conn,
                             const ::google::protobuf::Message& message) {
  Buffer buf;
  FillEmptyBuffer(&buf, message);
  conn->Send(&buf);
}

void ProtobufCodecLite::FillEmptyBuffer(
    Buffer* buf, const google::protobuf::Message& message) {
  assert(buf->ReadableBytes() == 0);
  buf->Append(tag_);

  int byte_size = SerializeToBuffer(message, buf);
  int32_t check_sum =
      Checksum(buf->Peek(), static_cast<int>(buf->ReadableBytes()));
  buf->AppendInt32(check_sum);
  assert(buf->ReadableBytes() == tag_.size() + byte_size + kChecksumLen);
  (void)byte_size;
  int32_t len =
      socket::HostToNetwork32(static_cast<int32_t>(buf->ReadableBytes()));
  buf->Prepend(&len, sizeof len);
}

void ProtobufCodecLite::OnMessage(const std::shared_ptr<TcpConnection>& conn,
                                  Buffer* buf) {
  while (buf->ReadableBytes() >=
         static_cast<uint32_t>(kMinMessageLen + kHeaderLen)) {
    const int32_t len = buf->PeekInt32();
    if (len > kMaxMessageLen || len < kMinMessageLen) {
      error_callback_(conn, buf, kInvalidLength);
      break;
    } else if (buf->ReadableBytes() >= kHeaderLen + len) {
      if (raw_callback_ &&
          !raw_callback_(conn,
                         std::string_view(buf->Peek(), kHeaderLen + len))) {
        buf->Retrieve(kHeaderLen + len);
        continue;
      }

      MessagePtr message(prototype_->New());
      ErrorCode error_code =
          Parse(buf->Peek() + kHeaderLen, len, message.get());
      if (error_code == kNoError) {
        message_callback_(conn, message);
        buf->Retrieve(kHeaderLen + len);
      } else {
        error_callback_(conn, buf, error_code);
        break;
      }
    } else {
      break;
    }
  }
}

bool ProtobufCodecLite::ParseFromBuffer(std::string_view buf,
                                        google::protobuf::Message* message) {
  return message->ParseFromArray(buf.data(), buf.size());
}

int ProtobufCodecLite::SerializeToBuffer(
    const google::protobuf::Message& message, Buffer* buf) {
   BufferOutputStream os(buf);
   message.SerializeToZeroCopyStream(&os);
   return static_cast<int>(os.ByteCount());

#if GOOGLE_PROTOBUF_VERSION > 3009002
  int byte_size = google::protobuf::internal::ToIntSize(message.ByteSizeLong());
#else
  int byte_size = message.ByteSize();
#endif
  buf->EnsureWritableBytes(byte_size + kChecksumLen);

  uint8_t* start = reinterpret_cast<uint8_t*>(buf->BeginWrite());
  uint8_t* end = message.SerializeWithCachedSizesToArray(start);
  if (end - start != byte_size) {
#if GOOGLE_PROTOBUF_VERSION > 3009002
    ByteSizeConsistencyError(
        byte_size,
        google::protobuf::internal::ToIntSize(message.ByteSizeLong()),
        static_cast<int>(end - start));
#else
    ByteSizeConsistencyError(byte_size, message.ByteSize(),
                             static_cast<int>(end - start));
#endif
  }
  buf->HasWritten(byte_size);
  return byte_size;
}

namespace {
const std::string kNoErrorStr = "NoError";
const std::string kInvalidLengthStr = "InvalidLength";
const std::string kCheckSumErrorStr = "CheckSumError";
const std::string kInvalidNameLenStr = "InvalidNameLen";
const std::string kUnknownMessageTypeStr = "UnknownMessageType";
const std::string kParseErrorStr = "ParseError";
const std::string kUnknownErrorStr = "UnknownError";
}  // namespace

const std::string& ProtobufCodecLite::ErrorCodeToString(ErrorCode error_code) {
  switch (error_code) {
    case kNoError:
      return kNoErrorStr;
    case kInvalidLength:
      return kInvalidLengthStr;
    case kCheckSumError:
      return kCheckSumErrorStr;
    case kInvalidNameLen:
      return kInvalidNameLenStr;
    case kUnknownMessageType:
      return kUnknownMessageTypeStr;
    case kParseError:
      return kParseErrorStr;
    default:
      return kUnknownErrorStr;
  }
}

void ProtobufCodecLite::DefaultErrorCallback(
    const std::shared_ptr<TcpConnection>& conn, Buffer* buf,
    ErrorCode error_code) {
  if (conn && conn->Connected()) {
    conn->Shutdown();
  }
}

int32_t ProtobufCodecLite::AsInt32(const char* buf) {
  int32_t be32 = 0;
  ::memcpy(&be32, buf, sizeof(be32));
  return socket::NetworkToHost32(be32);
}

int32_t ProtobufCodecLite::Checksum(const void* buf, int len) {
  return static_cast<int32_t>(
      ::adler32(1, static_cast<const Bytef*>(buf), len));
}

bool ProtobufCodecLite::ValidateChecksum(const char* buf, int len) {
  // check sum
  int32_t expected_checksum = AsInt32(buf + len - kChecksumLen);
  int32_t check_sum = Checksum(buf, len - kChecksumLen);
  return check_sum == expected_checksum;
}

ProtobufCodecLite::ErrorCode ProtobufCodecLite::Parse(
    const char* buf, int len, ::google::protobuf::Message* message) {
  ErrorCode error = kNoError;

  if (ValidateChecksum(buf, len)) {
    if (memcmp(buf, tag_.data(), tag_.size()) == 0) {
      // parse from buffer
      const char* data = buf + tag_.size();
      int32_t data_len = len - kChecksumLen - static_cast<int>(tag_.size());
      if (ParseFromBuffer(std::string_view(data, data_len), message)) {
        error = kNoError;
      } else {
        error = kParseError;
      }
    } else {
      error = kUnknownMessageType;
    }
  } else {
    error = kCheckSumError;
  }

  return error;
}
